
import pandas as pd
import numpy as np
from sklearn import tree
import matplotlib.pyplot as plt
import joblib
from utils.utils import get_classes



classes_path    = '../model_data/news_classes.txt'

#-------------------------------------#
#   设置数据集的路径
#   csv_path   训练集数据和标签
#   valid_csv_path   验证集数据和标签
#-------------------------------------#
train_csv_path   = "../train_Dataset.csv"
valid_csv_path   = "../valid_Dataset.csv"


if __name__ == "__main__":

    #----------------------#
    #   读取数据集对应的csv
    #----------------------#
    train_csv = pd.read_csv(train_csv_path)
    train_data = np.array(train_csv)

    valid_csv = pd.read_csv(valid_csv_path)
    valid_data = np.array(valid_csv)

    num_train = len(train_data)
    num_val = len(valid_data)

    print(f"训练样本数量:{num_train}, 验证样本数量:{num_val}")

    all_data = np.vstack((train_data, valid_data))

    train_x, train_y = all_data[:, :-1], all_data[:, -1].astype(int)

    print("DT")
    model = tree.DecisionTreeClassifier(max_depth=30)
    model.fit(train_x, train_y)
    train_score = model.score(train_x, train_y)
    print("训练集：", train_score)

    joblib.dump(model, 'all_samples_Dtree.pth')

    text_representation = tree.export_text(model)
    with open("all_decision_tree.log", "w") as fout:
        fout.write(text_representation)

    fig = plt.figure(figsize=(25, 20))

    classes, _ = get_classes(classes_path)

    _ = tree.plot_tree(
        model,
        feature_names=list(train_csv.columns[:-1]),
        class_names=classes,
        rounded=True,
        filled=True
    )
    fig.savefig("all_decision_tree.png")


